import os
from datetime import datetime
from src.common import commonUtils
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from langgraph.graph import StateGraph, START, END
from langchain_core.prompts import ChatPromptTemplate
from typing import List, TypedDict, Dict, Tuple, Union

logger = getLogger()

class Plan(BaseModel):
    index: int = Field(description = "计划序号")
    detail: str = Field(description = "计划详情描述")
    funcName: str = Field(description = "可用的工具名")
    funcParams: Dict[str, str] = Field(description = "json格式的工具的输入参数")

class PlanList(BaseModel):
    plans: List[Plan] = Field(description = "Plan对象的列表")

class Response(BaseModel):
    answer: str

class ReplanAct(BaseModel):
    action: Union[Response, PlanList] = Field(description = "要执行的操作。若已经得到正确答案，请使用 Response。若需进一步使用工具获取答案，请使用 PlanList。")

class PlanExecuteState(TypedDict):
    query: str
    plans: List[Plan]
    past_steps: List[Tuple]
    results: dict
    answer: str

class PlanExecuteAgent:

    def __init__(self, llm_model, agent_tools):
        self.llm_model = llm_model
        self.agent_tools = agent_tools

    def plan_node(self, state):
        logger.info("PlanExecuteAgent plan_node start")
        query = state["query"]
        template = """
            当前时间是：{current_date}

            针对以下任务，制定分步骤解决问题的计划。不要添加任何多余的步骤。
            针对每个计划，需明确指定用于获取证据的外部工具及其输入参数。
            可将证据存储到变量#E中，以便后续工具调用。
            （计划，#E1，计划，#E2，计划，...）

            工具可以是以下之一：{tool_descs}

            例如，
            任务：托马斯、托比和丽贝卡一周共工作了157小时。
            Thomas工作了x小时。
            托比的工作时间比托马斯少10小时，是托马斯工作时间的两倍，而丽贝卡比托比少8小时。
            丽贝卡工作了多少小时？

            计划：假设托马斯工作了x小时，将问题转化为代数表达式，并使用Wolfram Alpha求解。#E1=WolframAlpha[Solve x+（2x−10）+（（2x−10）−8）=157]
            计划：计算托马斯的工作时长。#E2=LLM[根据#E1求解x]
            计划：计算丽贝卡工作的时间。#E3=Calculator[（2∗#E2−10）−8]

            开始！请详细描述您的计划。每个计划只能跟随一个#E，并严格按格式输出计划。

            任务: {task}
        """
        tool_descs = commonUtils.build_tools_description(self.agent_tools)
        current_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        prompt = template.format(current_date = current_date, tool_descs = tool_descs, task = query)
        prompt = ChatPromptTemplate.from_messages([("user", prompt)])
        structure_llm = self.llm_model.with_structured_output(PlanList)
        response = (prompt | structure_llm).invoke({ "task": query})
        logger.info(f"PlanExecuteAgent plan_node response: {response}")
        return { "plans": response.plans }

    def execute_node(self, state):
        logger.info("PlanExecuteAgent execute_node start")
        plans = state["plans"]
        step = plans[0]
        logger.info(f"PlanExecuteAgent execute_node step: {step}")
        response = ""
        if step.funcName:
            for tool in self.agent_tools:
                if step.funcName == tool.name:
                    response = tool.func(**step.funcParams)
        else:
            plan_str = "\n".join(f"{i + 1}.{plan}" for i, plan in enumerate(plans))
            prompt = f"""
                你是一个智能AI助手，请推理并执行任务。
                对于以下计划：{plan_str}
                您需要执行步骤{1}，{step}。
            """
            response = self.llm_model.invoke(prompt)

        past_steps = state.get("past_steps", [])
        past_steps.append((step, response))
        return { "past_steps": past_steps }

    def replan_node(self, state):
        logger.info("PlanExecuteAgent replan_node start")
        replan_template ="""
            针对给定的目标，制定一个简单的分步计划。
            此计划应包含若干具体任务，若正确执行，即可得出正确答案。请勿添加任何多余步骤。
            最后一步的结果应为最终答案。确保每一步都包含所需信息，切勿跳过任何步骤。
                
            您的目标是：{query}
                
            你最初的计划是这样的：{plans}
                
            您当前已执行以下步骤：{past_steps}
                
            相应地更新您的计划。
            如果不需要更多步骤，且可以返回用户界面，则直接回复该信息。否则，请填写计划。
            只添加仍需完成的步骤。不要将以前完成的步骤作为计划的一部分返回。
        """
        replan_prompt = ChatPromptTemplate.from_template(replan_template)
        replan_chain = replan_prompt | self.llm_model.with_structured_output(ReplanAct)
        response = replan_chain.invoke(state)
        logger.info(f"PlanExecuteAgent replan_node response: {response}")
        if isinstance(response.action, Response):
            return { "answer": response.action.answer }
        else:
            return { "plans": response.action.plans }

    def route_node(self, state):
        logger.info("PlanExecuteAgent route_node start")
        answer = state.get("answer", None)
        logger.info(f"PlanExecuteAgent route_node answer: {answer}")
        if answer is not None:
            return END
        else:
            return "agent"

    def build_graph(self):
        graph = StateGraph(PlanExecuteState)
        graph.add_node("plan", self.plan_node)
        graph.add_node("agent", self.execute_node)
        graph.add_node("replan", self.replan_node)

        graph.add_edge(START, "plan")
        graph.add_edge("plan", "agent")
        graph.add_edge("agent", "replan")
        graph.add_conditional_edges("replan", self.route_node, { "agent": "agent" })

        workflow = graph.compile()

        save_path = "D:/Downloads/taixu/images/agentics"
        image_path = os.path.join(save_path, "Agent_PlanExecute_Workflow.png")
        if not os.path.exists(save_path):
            os.makedirs(save_path)

        with open(image_path, 'wb') as file:
            file.write(workflow.get_graph().draw_mermaid_png())

        return workflow

    def invoke(self, query):
        logger.info(f"PlanExecuteAgent invoke query: {query}")
        workflow = self.build_graph()
        response = workflow.invoke({ "query": query })
        answer = response.get("answer", None)
        logger.info(f"PlanExecuteAgent invoke answer len: {len(str(answer))}")
        return { "chain_result": answer }
